
#include <stdio.h>
#include <stdlib.h>
#include <math.h>
#include <apr_time.h>
#include "CongestionControl.h"

#define MINIMUM_ACK_PROPORTION 0.001

/** Constructor for the congestion control class.
 * 
 * @param ibytes_per_second
 * Maximum bandwidth for the associated data stream.
 * 
 * @param iusec_ack_timeout
 * The timeout before which an acknowledgement must be
 * received and passed back to this class.
 * 
 * @param icongestion_window_size
 * The window size giving the maximum number of packets
 * allowed in flight at any given time. */
CongestionControl::CongestionControl(int64_t ibytes_per_second,
   int iusec_ack_timeout, double icongestion_window_size) :
   RateControl(ibytes_per_second, icongestion_window_size)
{
   usec_ack_timeout = iusec_ack_timeout;
   packet_prop = 1.;
}

/** Checks for timed out packets and removes them from the congestion
 * window structure if they exist. */
void CongestionControl::checkForTimeouts()
{
   // Remember that the defaults for TCP aren't useful here - 
   // dropping packets is much more acceptable over UDP!
   double window_mul = 1.0;
   double ack_prop_mul = 1.0;
   // Check for timed out acknowledgements.
   int64_t thistime = apr_time_now();
   for (std::map<uint32_t, int64_t>::iterator
      i = congestion_window.begin(); i != congestion_window.end(); i++)
   {
      if ((thistime - (i->second)) > usec_ack_timeout)
      {
         int32_t elem = i->first;
         i = congestion_window.begin();
         congestion_window.erase(elem);
         window_mul *= 0.5;
         ack_prop_mul *= 2.0;
         last_reset_ts = thistime;
      }
   }
   congestion_window_size *= window_mul;
   congestion_window_size = (congestion_window_size < 2.) ? 2. : congestion_window_size;
   packet_prop *= ack_prop_mul;
   packet_prop = (packet_prop > 1.) ? 1. : packet_prop;
   packet_prop = (packet_prop < MINIMUM_ACK_PROPORTION) ? MINIMUM_ACK_PROPORTION : packet_prop;
}

/** Prepares the congestion control for sending a packet. Ensure that
 * bandwidth caps are met, and the congestion window is of the
 * appropriate size.
 * 
 * @param packet_id
 * The ID of the packet currently being sent.
 * 
 * @param datasize
 * The size of the packet to be sent.
 * 
 * @return
 * Either -1 for try again, the congestion window is full. Zero if
 * a receipt is not required by the protocol, one if it is.
 * Receipts (acks) are required by probabilistically selected packets.
 * Repeated calls will eventually go through, even if nothing is being
 * sent, because the congestion control will invalidate timed out data. */
int CongestionControl::prepareForSend(int32_t packet_id, int datasize)
{
   blockUntilAllowed(datasize);
   checkForTimeouts();
   if (congestion_window.size() >= congestion_window_size)
      return -1;
   // Now add the packet_id to the set
   int64_t timestamp = data_collection.back().first;
   congestion_window.insert(std::pair<uint32_t, int64_t>(packet_id, timestamp));
   int ask_for_receipt = ((rand() / (RAND_MAX + 1.)) < packet_prop);
   return ask_for_receipt;
}

/** Called when we want to process an acknowledgement from a remote
 * host that the given packet was received.
 * 
 * @param packet_id
 * A packet ID used to find the origin time of the packet.
 * 
 * @return
 * Zero if the packet ID was found in the database and deleted,
 * otherwise one if not found (no error really, probably timed out). */
int CongestionControl::acknowledgementReceived(int32_t packet_id)
{
   // Make a judgement on RTT to generate an increase in window size
   // / packet ack proportion
   std::map<uint32_t, int64_t>::iterator found = congestion_window.find(packet_id);
   if (found != congestion_window.end())
   {
      int64_t timenow = apr_time_now();
      int64_t send_timestamp = found->second;
      int64_t packet_rtt = timenow - send_timestamp;
      int64_t time_since = timenow - last_reset_ts;
      double rtts_since = time_since / ((double)packet_rtt);
      if (rtts_since > 0.)
      {
         congestion_window_size += (rtts_since * 1.0);
         packet_prop *= exp(-rtts_since * log(0.9));
      }
      packet_prop = (packet_prop < MINIMUM_ACK_PROPORTION) ? MINIMUM_ACK_PROPORTION : packet_prop;
      last_reset_ts = timenow;
   } else
      return 1;
   congestion_window.erase(packet_id);
   return 0;
}
